#include "ip.h"
#include "arp.h"
#include "icmp.h"
#include "udp.h"
#include <string.h>
#include <stdio.h>


static int ip_id = 0;

int check_ip_header(ip_hdr_t *ip_header)
{
    if(ip_header->version != IP_VERSION_4)
        return 0;
    if(ip_header->total_len < 20)
        return 0;
    if(ip_header->hdr_len * IP_HDR_LEN_PER_BYTE < 20)
        return 0;
    return 1;
}

/**
 * @brief 处理一个收到的数据包
 *        你首先需要做报头检查，检查项包括：版本号、总长度、首部长度等。
 * 
 *        接着，计算头部校验和，注意：需要先把头部校验和字段缓存起来，再将校验和字段清零，
 *        调用checksum16()函数计算头部检验和，比较计算的结果与之前缓存的校验和是否一致，
 *        如果不一致，则不处理该数据报。
 * 
 *        检查收到的数据包的目的IP地址是否为本机的IP地址，只处理目的IP为本机的数据报。
 * 
 *        检查IP报头的协议字段：
 *        如果是ICMP协议，则去掉IP头部，发送给ICMP协议层处理
 *        如果是UDP协议，则去掉IP头部，发送给UDP协议层处理
 *        如果是本实验中不支持的其他协议，则需要调用icmp_unreachable()函数回送一个ICMP协议不可达的报文。
 *          
 * @param buf 要处理的包
 */
void ip_in(buf_t *buf)
{
    // TODO 
    uint16_t header_checksum;
    ip_hdr_t *ip_header = (ip_hdr_t*)(buf->data);
    
    if(!check_ip_header(ip_header))
    {
        printf("IP header is wrong\n");
    }

    // check sum
    header_checksum = ip_header->hdr_checksum;
    ip_header->hdr_checksum = 0;
    if(checksum16((uint16_t*)ip_header, ip_header->hdr_len * IP_HDR_LEN_PER_BYTE / 2) != header_checksum)
        return;
    ip_header->hdr_checksum = header_checksum;

    // only process fitted ip
    if(memcmp(ip_header->dest_ip, net_if_ip, NET_IP_LEN) != 0)
        return;

    // determine how to process packet
    switch (ip_header->protocol)
    {
    case NET_PROTOCOL_ICMP:
        buf_remove_header(buf, ip_header->hdr_len * IP_HDR_LEN_PER_BYTE);
        icmp_in(buf, ip_header->src_ip);
        break;
    case NET_PROTOCOL_UDP:
        buf_remove_header(buf, ip_header->hdr_len * IP_HDR_LEN_PER_BYTE);
        udp_in(buf, ip_header->src_ip);
        break;
    default:
        icmp_unreachable(buf, ip_header->src_ip, ICMP_CODE_PROTOCOL_UNREACH);
        break;
    }
}

/**
 * @brief 处理一个要发送的分片
 *        你需要调用buf_add_header增加IP数据报头部缓存空间。
 *        填写IP数据报头部字段。
 *        将checksum字段填0，再调用checksum16()函数计算校验和，并将计算后的结果填写到checksum字段中。
 *        将封装后的IP数据报发送到arp层。
 * 
 * @param buf 要发送的分片
 * @param ip 目标ip地址
 * @param protocol 上层协议
 * @param id 数据包id
 * @param offset 分片offset，必须被8整除
 * @param mf 分片mf标志，是否有下一个分片
 */
void ip_fragment_out(buf_t *buf, uint8_t *ip, net_protocol_t protocol, int id, uint16_t offset, int mf)
{
    // TODO
    buf_add_header(buf, sizeof(ip_hdr_t));
    
    // fill in ip header
    ip_hdr_t *ip_header = (ip_hdr_t*)buf->data;
    ip_header->version = IP_VERSION_4;
    ip_header->hdr_len = sizeof(ip_hdr_t) / IP_HDR_LEN_PER_BYTE;
    ip_header->tos = 0;
    ip_header->total_len = swap16(buf->len);
    ip_header->id = swap16((uint16_t)id);
    ip_header->flags_fragment = swap16(mf<<8 | offset);
    ip_header->ttl = IP_DEFALUT_TTL;
    ip_header->protocol = (uint8_t)protocol;
    ip_header->hdr_checksum = 0;
    memcpy(ip_header->src_ip, net_if_ip, NET_IP_LEN);
    memcpy(ip_header->dest_ip, ip, NET_IP_LEN);
    
    // calculate checksum
    ip_header->hdr_checksum = checksum16((uint16_t*)ip_header, ip_header->hdr_len * IP_HDR_LEN_PER_BYTE / 2);

    arp_out(buf, ip, NET_PROTOCOL_IP);
}

/**
 * @brief 处理一个要发送的数据包
 *        你首先需要检查需要发送的IP数据报是否大于以太网帧的最大包长（1500字节 - ip包头长度）。
 *        
 *        如果超过，则需要分片发送。 
 *        分片步骤：
 *        （1）调用buf_init()函数初始化buf，长度为以太网帧的最大包长（1500字节 - ip包头头长度）
 *        （2）将数据报截断，每个截断后的包长度 = 以太网帧的最大包长，调用ip_fragment_out()函数发送出去
 *        （3）如果截断后最后的一个分片小于或等于以太网帧的最大包长，
 *             调用buf_init()函数初始化buf，长度为该分片大小，再调用ip_fragment_out()函数发送出去
 *             注意：最后一个分片的MF = 0
 *    
 *        如果没有超过以太网帧的最大包长，则直接调用调用ip_fragment_out()函数发送出去。
 * 
 * @param buf 要处理的包
 * @param ip 目标ip地址
 * @param protocol 上层协议
 */
void ip_out(buf_t *buf, uint8_t *ip, net_protocol_t protocol)
{
    // TODO 
    if(buf->len <= ETHERNET_MTU - sizeof(ip_hdr_t))
    {
        ip_fragment_out(buf, ip, protocol, ip_id++, 0, 0);
        return;
    }

    // if packet length is larger than MTU, fragment ip packet
    uint16_t i, len, offset;
    buf_t ip_buf;
    uint16_t d = ETHERNET_MTU - sizeof(ip_hdr_t);     // length of each packet
    d /= IP_HDR_OFFSET_PER_BYTE;
    d *= IP_HDR_OFFSET_PER_BYTE;
    
    uint16_t n = buf->len / d;          // totol size of packets
    if(buf->len % d != 0)
        n++;

    // 1 ~ n-1th fragments
    for(i=0; i<n-1; i++)
    {
        len = d;
        offset = d*i/IP_HDR_OFFSET_PER_BYTE;
        buf_init(&ip_buf, len);
        memcpy(ip_buf.data, (buf->data + offset*IP_HDR_OFFSET_PER_BYTE) , len);
        ip_fragment_out(&ip_buf, ip, protocol, ip_id, offset, IP_MORE_FRAGMENT);
    }
    // the last fragment
    len = (buf->len) - (n-1) * d;
    offset = d*i/IP_HDR_OFFSET_PER_BYTE;
    buf_init(&ip_buf, len);
    memcpy(ip_buf.data, (buf->data + offset*IP_HDR_OFFSET_PER_BYTE) , len);
    ip_fragment_out(&ip_buf, ip, protocol, ip_id++, offset, 0);
}
